/* compile with: gcc -Wall -O3 -o eigen-mult-timeseries eigen-mult-timeseries.c -lm -lgsl -lgslcblas */

/* ========================================================================= */
/* simulates the Eigen model with multiple mutations and frequency-dependent */
/* fitness; multiple populations                                             */
/* ========================================================================= */


/*  
    gnuplot not-so-faq im wiki
    
    in gnuplot heatmaps plotten mit
    set pm3d map
    splot 'file.dat' u(sing) 1:2:3 w(ith) pm3d
*/
#include <stdio.h>
#include <math.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_interp.h>

#define TYPE unsigned int
#define L 32
#define ONE 1
#define RRMASK 0x000000FF
#define MAXINT 0xFFFFFFFF
#define LRR 8

#define N 10000
#define TMAX 1000

#define NM 4 /* population number */

const TYPE masters[] = {0x00000000, 0xFFFFFF00, 0xAAAAAA00, 0x55555500, 
			0x11111100};

const double alpha=5.;
const double gam=NM*10.;

/* nach IIb */

gsl_rng *r;
gsl_interp_accel *acc;

inline TYPE mutate (TYPE x, double mu);
inline int hd (TYPE x, TYPE y); // hamming dist

void reproduce (TYPE *pop, double mu, unsigned int cdis[NM][(ONE << LRR)]);

TYPE compute_consensus (TYPE *pop);
void compute_cdis (TYPE *pop, unsigned int *cdis, TYPE master);

int main (int argc, char **argv) {

  int i,j,k,t,s;
  TYPE pop[N];
  double mu; //tation rate
  TYPE smax[NM]; // maximal populated sequence
  unsigned int cdis[NM][(ONE << LRR)];

  // Random numbers
  const gsl_rng_type *T;
  gsl_rng_env_setup();
  acc=gsl_interp_accel_alloc();
  T=gsl_rng_gfsr4;
  r=gsl_rng_alloc (T);

  if (argc > 2) {
    mu=atof(argv[1]);
    s=atoi(argv[2]);
    fprintf(stderr, "mu=% .3e, s=%i\n", mu, s);
  } else {
    fprintf(stderr, "no mutation rate given!\n");
    exit(1);
  }

  gsl_rng_set (r, s);

  // pop initialization
  for (j=0; j<NM; j++) 
    for (i=j*N/NM; i<(j+1)*N/NM; i++) 
      pop[i]=(masters[j] & ~RRMASK) | gsl_rng_uniform_int(r, ONE << LRR);  
  /*      pop[i]=masters[j]; */

  // frequency of master sequences
  for (j=0; j<NM; j++)
    compute_cdis(pop, cdis[j], masters[j]);

  for (t=0; t<TMAX; t++) {
    
    reproduce (pop, mu, cdis);
    for (j=0; j<NM; j++)
      compute_cdis(pop, cdis[j], masters[j]);

    for (k=0; k<(ONE << LRR); k++) {
      printf("%i %i ", t, k);
      for (j=0; j<NM; j++)
	printf("%i ", cdis[j][k]);
      printf("\n");
    }
    printf("\n");
  } 

  for (j=0; j<NM; j++) {
    smax[j]=0;
    for (k=1; k<(ONE << LRR); k++)
      if (cdis[j][k] > cdis[j][smax[j]]) smax[j]=k;

    fprintf(stderr, "smax[%i]= %3i (%02X), neighbors: ", j, smax[j], smax[j]);
    for (k=0; k<LRR; k++)
      fprintf(stderr, "%i ", smax[j]^(ONE << k));
    fprintf(stderr, "\n");
  }

  

  return 0;
}

void compute_cdis (TYPE *pop, unsigned int *cdis, TYPE master) {

  int i,k;

  for (k=0; k<(ONE << LRR); k++)
    cdis[k]=0;
  
  for (i=0; i<N; i++)
    if ((pop[i] & ~RRMASK)==(master & ~RRMASK))
      cdis[(unsigned int)(pop[i] & RRMASK)]++;

  return;
}

void reproduce (TYPE* pop, double mu, unsigned int cdis[NM][(ONE << LRR)]) {

  int i,j,index; // i seque, j populations
  TYPE popnew[N];
  double ptot;
  double sel[N];
  double mut[N];
  double select[2*N];
  double delta[NM]; // wieviele in pop j den genotyp von typ i haben (in RR) 

  for (i=0; i<N; i++) {

    for (j=0; j<NM; j++) 
      delta[j]=cdis[j][(pop[i]) & RRMASK]/(double)N;
    
    mut[i]=1.;
    sel[i]=1; /* 1./DT */

    for (j=0; j<NM; j++)
      if ((pop[i] & ~RRMASK)==(masters[j] & ~RRMASK))
	mut[i]=alpha+gam*delta[(j+1)%NM];
  }
  
  /* sum of probabilities for normalization */
  ptot=0;
  for (i=0; i<N; i++) ptot+=sel[i]+mut[i];

  /* evaluate probabilities to be selected 
     (even index) and to be selected+mutated 
     (odd indices) */
  select[0]=sel[0]/ptot;
  select[1]=mut[0]/ptot+select[0];
  for (i=1; i<N; i++) {
    select[2*i]=sel[i]/ptot+select[2*i-1];
    select[2*i+1]=mut[i]/ptot+select[2*i];
  }
  /* get N random numbers in [0,1) */
  for (i=0; i<N; i++) {

    /* look up entry in select array */
    index=gsl_interp_bsearch (select, gsl_rng_uniform(r), 0, 2*N-2)+1;
    
    /* if index is odd, mutate this sequence with 
       probability mu, otherwise copy  */ 
          popnew[i]=(index%2==1 ? mutate(pop[index/2], mu) : pop[index/2]); 
    
  }    

  for (i=0; i<N; i++) pop[i]=popnew[i];

  return;
}


TYPE compute_consensus (TYPE *pop) {

  int i,k;
  int aseq[L];
  TYPE consensus;

  for (k=0; k<L; k++) aseq[k]=0;
				/* for every sequence in population */
  for (i=0; i<N; i++)
				/* look at bit k */
    for (k=0; k<L; k++) 
				/* if this bit is 1, increase counter */
      if (pop[i] & (ONE << k)) aseq[k]++;
  
  consensus=gsl_rng_uniform_int(r, MAXINT);
  for (k=0; k<L; k++)
    if (aseq[k] >= N/2) consensus |= (ONE << k);
    else consensus &= ~(ONE << k);

  return consensus;
}

inline TYPE mutate (TYPE x, double mu) {

  int i;
  TYPE out=x;

  for (i=0; i<L; i++)
    if (gsl_rng_uniform(r) < mu)
      out^=(ONE << i);
  
  return out;
}

inline int hd (TYPE x, TYPE y) {

  int dist=0;
  TYPE val=x ^ y;

  while (val) {

    ++dist;
    val &= val-1;
  }

  return dist;
}
